import importlib
import pkgutil
from pathlib import Path
from typing import Dict, Callable, List

# 全局仓库：{装饰器名: [(类, 方法), ...]}
_METHOD_REGISTRY: Dict[str, List[Callable]] = {}

def register(func_or_name: str | Callable):
    """
    装饰器：可带括号或不带括号
    @register
    @register("custom_name")
    """
    if isinstance(func_or_name, str):
        name = func_or_name
        def decorator(fn):
            _METHOD_REGISTRY.setdefault(name, []).append(fn)
            return fn
        return decorator
    else:
        _METHOD_REGISTRY.setdefault("default", []).append(func_or_name)
        return func_or_name

def scan_directory(package_path: Path, package_name: str):
    """扫描 package_path 目录下所有模块，并 import 它们"""
    for _, module_name, _ in pkgutil.iter_modules([str(package_path)]):
        full_name = f"{package_name}.{module_name}"
        importlib.import_module(full_name)

def collect(name: str = "default") -> List[Callable]:
    """取出被 @register 标记的方法"""
    return _METHOD_REGISTRY.get(name, [])